package org.tao.lightningmq.broker.rebalance;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.tao.lightningmq.broker.cache.CommonCache;
import org.tao.lightningmq.broker.model.MqTopicModel;
import org.tao.lightningmq.broker.rebalance.strategy.IReBalanceStrategy;
import org.tao.lightningmq.broker.rebalance.strategy.RandomReBalanceStrategyImpl;
import org.tao.lightningmq.broker.rebalance.strategy.ReBalanceInfo;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * @Author lht
 * @date 2025/1/19 - 23:04
 * @description:
 */
public class ConsumerInstancePool {

    private static final Logger LOGGER = LoggerFactory.getLogger(ConsumerInstancePool.class);

    private Map<String, List<ConsumerInstance>> consumeInstanceMap = new ConcurrentHashMap<>();
    private static Map<String, IReBalanceStrategy> reBalanceStrategyMap = new HashMap<>();
    private ReBalanceInfo reBalanceInfo = new ReBalanceInfo();

    static {
        reBalanceStrategyMap.put("random", new RandomReBalanceStrategyImpl());
        reBalanceStrategyMap.put("range", new RandomReBalanceStrategyImpl());
    }

    /**
     * 将消费者实例加入到池中
     *
     * @param consumerInstance 消费者实例
     */
    public void addInstancePool(ConsumerInstance consumerInstance) {
        String topic = consumerInstance.getTopic();
        MqTopicModel mqTopicModel = CommonCache.getMqTopicModelMap().get(topic);
        Objects.requireNonNull(mqTopicModel, "topic非法");
        synchronized (this) {
            // 使用链表方便删除
            List<ConsumerInstance> consumerInstances = consumeInstanceMap
                    .computeIfAbsent(topic, k -> new LinkedList<>());
            for (ConsumerInstance instance : consumerInstances) {
                if (instance.getConsumerReqId().equals(consumerInstance.getConsumerReqId())) {
                    return;
                }
            }
            consumerInstances.add(consumerInstance);
            Set<String> consumerGroupSet = reBalanceInfo.getChangeConsumerGroupMap()
                    .computeIfAbsent(topic, k -> new HashSet<>());
            consumerGroupSet.add(consumerInstance.getConsumeGroup());
            doReBalance();
            LOGGER.info("new instance add in pool:{}", consumerInstance);
        }
    }

    public synchronized void removeFromInstancePool(String reqId) {
        consumeInstanceMap.forEach((k, v) -> {
            v.removeIf(consumerInstance -> consumerInstance.getConsumerReqId().equals(reqId));
        });
        Map<String, Map<String, List<ConsumerInstance>>> consumeHoldMap = CommonCache.getConsumeHoldMap();
        consumeHoldMap.values().stream()
                .flatMap(map -> map.values().stream())
                .forEach(consumers -> consumers.removeIf(consumerInstance -> consumerInstance.getConsumerReqId().equals(reqId)));
    }

    public void startReBalanceJob() {
        new Thread(() -> {
            while (true) {
                try {
                    TimeUnit.SECONDS.sleep(10);
                    doReBalance();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }, "reBalance-job").start();
    }

    private synchronized void doReBalance() {
        // FIXME 在存在多个主题时有bug
        String reBalanceStrategy = CommonCache.getGlobalProperties().getReBalanceStrategy();
        reBalanceInfo.setConsumeInstanceMap(consumeInstanceMap);
        reBalanceStrategyMap.get(reBalanceStrategy).doReBalance(reBalanceInfo);
        reBalanceInfo.getChangeConsumerGroupMap().clear();
        LOGGER.info("do reBalance,{}", CommonCache.getConsumeHoldMap());
    }

}
